package org.apache.zeppelin.presto;

import io.prestosql.jdbc.PrestoConnection;
import io.prestosql.jdbc.PrestoResultSet;
import io.prestosql.jdbc.PrestoStatement;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.exception.ExceptionUtils;
import org.apache.zeppelin.interpreter.*;
import org.apache.zeppelin.interpreter.util.SqlSplitter;
import org.apache.zeppelin.tabledata.TableDataUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicReference;

import static org.apache.commons.lang3.StringUtils.containsIgnoreCase;

/**
 * @author leojie 2022/1/16 7:57 下午
 */
public class PrestoInterpreter extends AbstractInterpreter {
    private static final Logger LOGGER = LoggerFactory.getLogger(PrestoInterpreter.class);

    static final String COMMON_KEY = "common";
    static final String MAX_LINE_KEY = "max_count";
    private static final String PRESTO_JDBC_URL = "presto.jdbc.url";
    private static final String PRESTO_JDBC_USER = "presto.jdbc.user";

    private PrestoConnection connection = null;
    private PrestoStatement statement = null;

    private final int maxLineResults;
    private int maxRows;

    private SqlSplitter sqlSplitter;
    private final Map<String, Boolean> isFirstRefreshMap = new HashMap<>();

    static final String EMPTY_COLUMN_VALUE = "";

    private static final char WHITESPACE = ' ';
    private static final char NEWLINE = '\n';
    private static final char TAB = '\t';
    private static final String TABLE_MAGIC_TAG = "%table ";
    private static final String EXPLAIN_PREDICATE = "EXPLAIN ";

    private static final LinkedBlockingQueue<Double> PROGRESS_QUEUE = new LinkedBlockingQueue<>();
    private static final LinkedBlockingQueue<String> PROGRESS_STATE_QUEUE = new LinkedBlockingQueue<>();

    ExecutorService progressExecutorService = Executors.newFixedThreadPool(2);

    public PrestoInterpreter(Properties properties) {
        super(properties);
        maxLineResults = 1000;
    }

    @Override
    public ZeppelinContext getZeppelinContext() {
        return null;
    }

    private InterpreterResult executeSql(String sql, InterpreterContext context) throws InterpreterException {
        if (this.connection == null) {
            Properties properties = new Properties();
            properties.setProperty("user", getProperty(PRESTO_JDBC_USER));
            try {
                this.connection = DriverManager.getConnection(getProperty(PRESTO_JDBC_URL), properties).unwrap(PrestoConnection.class);
            } catch (SQLException e) {
                LOGGER.error("Created presto connection failed, the reason is {}", e.getMessage());
            }
        }
        if (this.statement == null) {
            try {
                this.statement = this.connection.createStatement().unwrap(PrestoStatement.class);
            } catch (SQLException e) {
                LOGGER.error("Created presto statement failed, the reason is {}", e.getMessage());
            }
            LOGGER.info("Created presto connection successfully.");
        }

        ResultSet resultSet = null;
        // String paragraphId = context.getParagraphId();

        try {
            List<String> sqlArray = sqlSplitter.splitSql(sql);
            for (String sqlToExecute : sqlArray) {
                String sqlTrimmedLowerCase = sqlToExecute.trim().toLowerCase();
                if (sqlTrimmedLowerCase.startsWith("set ") ||
                        sqlTrimmedLowerCase.startsWith("list ") ||
                        sqlTrimmedLowerCase.startsWith("add ") ||
                        sqlTrimmedLowerCase.startsWith("delete ")) {
                    sqlToExecute = sqlToExecute.trim();
                }
                LOGGER.info("Execute sql: " + sqlToExecute);

                statement.setFetchSize(context.getIntLocalProperty("limit", getMaxResult()));
                statement.setLargeMaxRows(context.getIntLocalProperty("limit", maxRows));

                try {
                    boolean isResultSetAvailable = statement.execute(sqlToExecute);

                    AtomicReference<String> queryId = new AtomicReference<>();
                    statement.setProgressMonitor(queryStats -> {

                        progressExecutorService.execute(() -> {
                            final double currentProgress = queryStats.getProgressPercentage().orElse(0);
                            try {
                                PROGRESS_QUEUE.put(currentProgress);
                                PROGRESS_STATE_QUEUE.put(queryStats.getState());
                                System.out.println("当前查询[" + queryId.get() + "]的状态：" + PROGRESS_STATE_QUEUE.take() + " ,执行进度：" + PROGRESS_QUEUE.take());
                            }catch (InterruptedException ex) {
                                ex.printStackTrace();
                            }
                        });
                        queryId.set(queryStats.getQueryId());
                    });


                    if (isResultSetAvailable) {
                        resultSet = statement.getResultSet().unwrap(PrestoResultSet.class);

                        // Regards that the command is DDL.
                        if (isDDLCommand(statement.getUpdateCount(),
                                resultSet.getMetaData().getColumnCount())) {
                            context.out.write("%text Query executed successfully.\n");
                        } else {
                            String template = context.getLocalProperties().get("template");
                            if (!StringUtils.isBlank(template)) {
                                resultSet.next();
                                SingleRowInterpreterResult singleRowResult =
                                        new SingleRowInterpreterResult(getFirstRow(resultSet), template, context);

                                if (isFirstRefreshMap.get(context.getParagraphId())) {
                                    context.out.write(singleRowResult.toAngular());
                                    context.out.write("\n%text ");
                                    context.out.flush();
                                    isFirstRefreshMap.put(context.getParagraphId(), false);
                                }
                                singleRowResult.pushAngularObjects();

                            } else {
                                String results = getResults(resultSet,
                                        !containsIgnoreCase(sqlToExecute, EXPLAIN_PREDICATE));
                                context.out.write(results);
                                context.out.write("\n%text ");
                                context.out.flush();
                            }
                        }
                    } else {
                        // Response contains either an update count or there are no results.
                        int updateCount = statement.getUpdateCount();
                        context.out.write("\n%text " +
                                "Query executed successfully. Affected rows : " +
                                updateCount + "\n");
                    }
                } finally {
                    if (resultSet != null) {
                        try {
                            resultSet.close();
                        } catch (SQLException e) { /*ignored*/ }
                    }
                    try {
                        statement.close();
                    } catch (SQLException e) { /*ignored*/ }

                }
            }
        } catch (Throwable e) {
            LOGGER.error("Cannot run " + sql, e);
            if (e instanceof SQLException) {
                return new InterpreterResult(InterpreterResult.Code.ERROR, e.getMessage());
            } else {
                return new InterpreterResult(InterpreterResult.Code.ERROR, ExceptionUtils.getStackTrace(e));
            }
        } finally {
            if (this.connection != null) {
                try {
                    this.connection.close();
                } catch (SQLException e) {
                    LOGGER.error("Closed presto connection failed. the reason is {}", e.getMessage());
                }
                this.connection = null;
            }
        }

        return new InterpreterResult(InterpreterResult.Code.SUCCESS);
    }

    @Override
    protected InterpreterResult internalInterpret(String cmd, InterpreterContext context) throws InterpreterException {
         return executeSql(cmd, context);
    }

    @Override
    public void open() throws InterpreterException {
        this.sqlSplitter = new SqlSplitter();
    }

    @Override
    public void close() throws InterpreterException {
    }

    @Override
    public void cancel(InterpreterContext context) throws InterpreterException {
        try {
            this.statement.cancel();
        } catch (SQLException e) {
            LOGGER.error("Cancel presto query failed, the reason is {}", e.getMessage());
        }
        LOGGER.info("Cancel presto query successfully.");
    }

    @Override
    public FormType getFormType() throws InterpreterException {
        return FormType.SIMPLE;
    }

    @Override
    public int getProgress(InterpreterContext context) throws InterpreterException {
        try {
            Double take = PROGRESS_QUEUE.take();
            return take.intValue();
        } catch (InterruptedException e) {
            e.printStackTrace();
            return 0;
        }
    }

    public int getMaxResult() {
        return maxLineResults;
    }

    private boolean isDDLCommand(int updatedCount, int columnCount) throws SQLException {
        return updatedCount < 0 && columnCount <= 0;
    }

    private List getFirstRow(ResultSet rs) throws SQLException {
        List list = new ArrayList();
        ResultSetMetaData md = rs.getMetaData();
        for (int i = 1; i <= md.getColumnCount(); ++i) {
            Object columnObject = rs.getObject(i);
            String columnValue = null;
            if (columnObject == null) {
                columnValue = "null";
            } else {
                columnValue = rs.getString(i);
            }
            list.add(columnValue);
        }
        return list;
    }

    private String getResults(ResultSet resultSet, boolean isTableType)
            throws SQLException {

        ResultSetMetaData md = resultSet.getMetaData();
        StringBuilder msg;
        if (isTableType) {
            msg = new StringBuilder(TABLE_MAGIC_TAG);
        } else {
            msg = new StringBuilder();
        }

        for (int i = 1; i < md.getColumnCount() + 1; i++) {
            if (i > 1) {
                msg.append(TAB);
            }
            if (StringUtils.isNotEmpty(md.getColumnLabel(i))) {
                msg.append(removeTablePrefix(replaceReservedChars(
                        TableDataUtils.normalizeColumn(md.getColumnLabel(i)))));
            } else {
                msg.append(removeTablePrefix(replaceReservedChars(
                        TableDataUtils.normalizeColumn(md.getColumnName(i)))));
            }
        }
        msg.append(NEWLINE);

        int displayRowCount = 0;
        boolean truncate = false;
        while (resultSet.next()) {
            if (displayRowCount >= getMaxResult()) {
                truncate = true;
                break;
            }
            for (int i = 1; i < md.getColumnCount() + 1; i++) {
                Object resultObject;
                String resultValue;
                resultObject = resultSet.getObject(i);
                if (resultObject == null) {
                    resultValue = "null";
                } else {
                    resultValue = resultSet.getString(i);
                }
                msg.append(replaceReservedChars(TableDataUtils.normalizeColumn(resultValue)));
                if (i != md.getColumnCount()) {
                    msg.append(TAB);
                }
            }
            msg.append(NEWLINE);
            displayRowCount++;
        }

        if (truncate) {
            msg.append("\n" + ResultMessages.getExceedsLimitRowsMessage(getMaxResult(),
                    String.format("%s.%s", COMMON_KEY, MAX_LINE_KEY)).toString());
        }
        return msg.toString();
    }

    private String removeTablePrefix(String columnName) {
        int index = columnName.indexOf(".");
        if (index > 0) {
            return columnName.substring(index + 1);
        } else {
            return columnName;
        }
    }

    private String replaceReservedChars(String str) {
        if (str == null) {
            return EMPTY_COLUMN_VALUE;
        }
        return str.replace(TAB, WHITESPACE).replace(NEWLINE, WHITESPACE);
    }

}
